"""
特征工程模块

负责提取和处理用于模型训练的特征，包括:
- 特征预处理与筛选
- 特征归一化
- 特征选择
- 滞后特征创建
- 目标变量生成
"""

import os
import logging
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Union, Tuple, Any, Set
from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler
from sklearn.feature_selection import SelectKBest, f_regression, mutual_info_regression
from scipy.stats import pearsonr

# 配置日志
logger = logging.getLogger(__name__)


class FeatureEngineer:
    """特征工程类，负责特征处理和转换"""
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化特征工程器
        
        Args:
            config: 配置字典，包含特征工程相关参数
        """
        self.config = config
        self.normalizers = {
            'standard': StandardScaler(),
            'minmax': MinMaxScaler(feature_range=(0, 1)),
            'robust': RobustScaler()
        }
        
        logger.info("初始化特征工程器")
    
    def preprocess_features(self, integrated_data: Dict[str, Dict[str, Dict[str, Any]]], feature_list: Optional[List[str]] = None) -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        预处理特征
        
        Args:
            integrated_data: 整合后的数据，键为股票代码，值为日期到特征的映射
            feature_list: 要保留的特征列表，如果为None则保留所有特征
            
        Returns:
            预处理后的特征数据，键为股票代码，值为日期到处理后特征的映射
        """
        logger.info("开始预处理特征")
        
        # 如果未指定特征列表，获取所有可用特征
        if feature_list is None:
            feature_list = self._get_all_feature_names(integrated_data)
            logger.info(f"未指定特征列表，将使用所有可用特征: {len(feature_list)}个")
        else:
            logger.info(f"使用指定的特征列表: {len(feature_list)}个")
        
        processed_data = {}
        for code, data in integrated_data.items():
            # 构建该股票的处理后数据
            processed_data[code] = {}
            
            # 遍历每个日期的数据
            for date, daily_data in data.items():
                # 筛选指定的特征
                processed_data[code][date] = {
                    feature: value for feature, value in daily_data.items() 
                    if feature in feature_list
                }
            
            logger.debug(f"预处理股票 {code} 的特征完成，保留了 {len(feature_list)} 个特征")
        
        logger.info(f"特征预处理完成，处理了 {len(integrated_data)} 只股票的数据")
        return processed_data
    
    def normalize_features(self, processed_data: Dict[str, Dict[str, Dict[str, float]]], method: str = 'standard') -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        特征归一化
        
        Args:
            processed_data: 预处理后的特征数据，键为股票代码，值为日期到特征的映射
            method: 归一化方法，可选值为'standard'(标准化)、'minmax'(最小最大缩放)、'robust'(稳健缩放)
            
        Returns:
            归一化后的特征数据，键为股票代码，值为日期到归一化特征的映射
        """
        logger.info(f"开始使用 {method} 方法进行特征归一化")
        
        # 获取归一化器，默认使用标准化
        normalizer = self.normalizers.get(method, self.normalizers['standard'])
        normalized_data = {}
        
        for code, data in processed_data.items():
            # 将字典数据转换为DataFrame以便进行归一化
            dates = list(data.keys())
            if not dates:
                logger.warning(f"股票 {code} 没有数据，跳过归一化")
                normalized_data[code] = {}
                continue
            
            # 提取所有特征名
            feature_names = list(data[dates[0]].keys())
            if not feature_names:
                logger.warning(f"股票 {code} 没有特征，跳过归一化")
                normalized_data[code] = data  # 返回原始数据
                continue
            
            # 构建特征矩阵
            feature_matrix = []
            for date in dates:
                feature_vector = [data[date].get(feature, np.nan) for feature in feature_names]
                feature_matrix.append(feature_vector)
            
            # 转换为NumPy数组
            feature_array = np.array(feature_matrix)
            
            # 对每个特征进行归一化
            normalized_array = np.zeros_like(feature_array)
            for i, feature in enumerate(feature_names):
                # 提取该特征的所有值
                feature_values = feature_array[:, i].reshape(-1, 1)
                
                # 检查是否有NaN值
                nan_mask = np.isnan(feature_values)
                non_nan_indices = ~nan_mask.flatten()
                
                if np.any(non_nan_indices):
                    # 只对非NaN值进行归一化
                    non_nan_values = feature_values[non_nan_indices]
                    normalized_values = normalizer.fit_transform(non_nan_values)
                    
                    # 将归一化后的值放回原位置
                    normalized_array[non_nan_indices, i] = normalized_values.flatten()
                    normalized_array[nan_mask.flatten(), i] = np.nan
                else:
                    # 如果全是NaN，则保持NaN
                    normalized_array[:, i] = np.nan
            
            # 将归一化后的数据转换回字典格式
            normalized_data[code] = {}
            for i, date in enumerate(dates):
                normalized_data[code][date] = {
                    feature: normalized_array[i, j]
                    for j, feature in enumerate(feature_names)
                }
            
            logger.debug(f"归一化股票 {code} 的特征完成，处理了 {len(feature_names)} 个特征")
        
        logger.info(f"特征归一化完成，处理了 {len(processed_data)} 只股票的数据")
        return normalized_data
    
    def select_features(self, normalized_data: Dict[str, Dict[str, Dict[str, float]]], method: str = 'correlation', n_features: int = 20, target_feature: str = 'close') -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        特征选择
        
        Args:
            normalized_data: 归一化后的特征数据，键为股票代码，值为日期到特征的映射
            method: 特征选择方法，可选值为'correlation'(相关性)、'mutual_info'(互信息)、'f_regression'(F检验)
            n_features: 选择的特征数量
            target_feature: 目标特征，用于计算相关性或互信息
            
        Returns:
            特征选择后的数据，键为股票代码，值为日期到选择后特征的映射
        """
        logger.info(f"开始使用 {method} 方法进行特征选择，保留 {n_features} 个特征")
        
        selected_features_data = {}
        all_selected_features = set()  # 所有股票选择的特征的并集
        
        for code, data in normalized_data.items():
            # 将字典数据转换为DataFrame
            df = pd.DataFrame.from_dict(data, orient='index')
            
            # 如果数据为空或目标特征不存在，则跳过该股票
            if df.empty or target_feature not in df.columns:
                logger.warning(f"股票 {code} 数据为空或不包含目标特征 {target_feature}，跳过特征选择")
                selected_features_data[code] = data  # 返回原始数据
                continue
            
            # 提取目标特征和其他特征
            y = df[target_feature].values
            X = df.drop(columns=[target_feature])
            feature_names = X.columns.tolist()
            
            # 特征选择
            selected_indices = []
            selected_features = []
            
            if method == 'correlation':
                # 使用相关性进行特征选择
                correlations = []
                for feature in feature_names:
                    # 计算与目标特征的相关性
                    mask = ~(np.isnan(X[feature]) | np.isnan(y))
                    if np.sum(mask) > 1:  # 需要至少两个有效值才能计算相关性
                        corr, _ = pearsonr(X[feature][mask], y[mask])
                        correlations.append((feature, abs(corr) if not np.isnan(corr) else 0))
                    else:
                        correlations.append((feature, 0))
                
                # 按相关性大小排序
                correlations.sort(key=lambda x: x[1], reverse=True)
                
                # 选择相关性最高的n_features个特征
                selected_features = [feature for feature, _ in correlations[:n_features]]
                
            elif method == 'mutual_info':
                # 使用互信息进行特征选择
                X_array = X.values
                selector = SelectKBest(mutual_info_regression, k=min(n_features, X.shape[1]))
                
                # 处理缺失值
                X_array_filled = np.nan_to_num(X_array, nan=0)
                selector.fit(X_array_filled, y)
                
                # 获取选择的特征索引
                selected_indices = selector.get_support(indices=True)
                selected_features = [feature_names[i] for i in selected_indices]
                
            elif method == 'f_regression':
                # 使用F检验进行特征选择
                X_array = X.values
                selector = SelectKBest(f_regression, k=min(n_features, X.shape[1]))
                
                # 处理缺失值
                X_array_filled = np.nan_to_num(X_array, nan=0)
                selector.fit(X_array_filled, y)
                
                # 获取选择的特征索引
                selected_indices = selector.get_support(indices=True)
                selected_features = [feature_names[i] for i in selected_indices]
                
            else:
                logger.warning(f"未知的特征选择方法: {method}，将使用所有特征")
                selected_features = feature_names
            
            # 确保目标特征也被选中
            if target_feature not in selected_features:
                selected_features.append(target_feature)
            
            # 更新所有选择的特征集合
            all_selected_features.update(selected_features)
            
            # 选择特征
            selected_features_data[code] = {}
            for date, daily_data in data.items():
                selected_features_data[code][date] = {
                    feature: value for feature, value in daily_data.items() 
                    if feature in selected_features
                }
            
            logger.debug(f"为股票 {code} 选择了 {len(selected_features)} 个特征: {', '.join(selected_features[:5])}...")
        
        logger.info(f"特征选择完成，处理了 {len(normalized_data)} 只股票的数据，总共选择了 {len(all_selected_features)} 个不同特征")
        return selected_features_data
    
    def create_lagged_features(self, data: Dict[str, Dict[str, Dict[str, float]]], lag_periods: List[int] = [1, 3, 5, 7]) -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        创建滞后特征
        
        Args:
            data: 特征数据，键为股票代码，值为日期到特征的映射
            lag_periods: 滞后周期列表，例如[1, 3, 5]表示创建1天前、3天前和5天前的滞后特征
            
        Returns:
            包含滞后特征的数据，键为股票代码，值为日期到特征的映射
        """
        logger.info(f"开始创建滞后特征，滞后周期: {lag_periods}")
        
        data_with_lagged_features = {}
        
        for code, stock_data in data.items():
            # 将字典数据转换为DataFrame以便处理时间序列
            df = pd.DataFrame.from_dict(stock_data, orient='index')
            
            # 确保索引是日期类型并排序
            df.index = pd.to_datetime(df.index)
            df = df.sort_index()
            
            # 获取原始特征名
            original_features = df.columns.tolist()
            
            # 创建滞后特征
            for lag in lag_periods:
                for feature in original_features:
                    # 创建滞后特征
                    df[f"{feature}_lag{lag}"] = df[feature].shift(lag)
            
            # 将DataFrame转回字典格式
            data_with_lagged_features[code] = df.to_dict(orient='index')
            
            # 将日期索引转回字符串格式
            data_with_lagged_features[code] = {
                date.strftime('%Y-%m-%d'): features 
                for date, features in data_with_lagged_features[code].items()
            }
            
            logger.debug(f"为股票 {code} 创建了滞后特征，从 {len(original_features)} 个原始特征扩展到 {len(df.columns)} 个特征")
        
        logger.info(f"滞后特征创建完成，处理了 {len(data)} 只股票的数据")
        return data_with_lagged_features
    
    def create_target_variable(self, data: Dict[str, Dict[str, Dict[str, float]]], prediction_horizon: int = 5, threshold: float = 0.02) -> Dict[str, Dict[str, Dict[str, Any]]]:
        """
        创建目标变量
        
        Args:
            data: 特征数据，键为股票代码，值为日期到特征的映射
            prediction_horizon: 预测周期，即未来多少天的价格变化
            threshold: 价格变化阈值，高于此阈值视为上涨，低于负阈值视为下跌，否则视为持平
            
        Returns:
            包含目标变量的数据，键为股票代码，值为日期到特征和目标变量的映射
        """
        logger.info(f"开始创建目标变量，预测周期: {prediction_horizon}天，价格变化阈值: {threshold}")
        
        data_with_target = {}
        
        for code, stock_data in data.items():
            # 将字典数据转换为DataFrame以便处理时间序列
            df = pd.DataFrame.from_dict(stock_data, orient='index')
            
            # 确保索引是日期类型并排序
            df.index = pd.to_datetime(df.index)
            df = df.sort_index()
            
            # 如果没有'close'列，则跳过
            if 'close' not in df.columns:
                logger.warning(f"股票 {code} 没有'close'列，无法创建目标变量")
                data_with_target[code] = stock_data
                continue
            
            # 计算未来价格变化
            df['future_price'] = df['close'].shift(-prediction_horizon)
            df['price_change'] = (df['future_price'] - df['close']) / df['close']
            
            # 创建目标变量（分类）
            df['target_class'] = np.where(df['price_change'] > threshold, 1,  # 上涨
                                 np.where(df['price_change'] < -threshold, -1,  # 下跌
                                 0))  # 持平
            
            # 创建目标变量（回归）
            df['target_regression'] = df['price_change']
            
            # 删除辅助列
            df = df.drop(columns=['future_price'])
            
            # 将DataFrame转回字典格式
            data_with_target[code] = df.to_dict(orient='index')
            
            # 将日期索引转回字符串格式
            data_with_target[code] = {
                date.strftime('%Y-%m-%d'): features 
                for date, features in data_with_target[code].items()
            }
            
            logger.debug(f"为股票 {code} 创建了目标变量")
        
        logger.info(f"目标变量创建完成，处理了 {len(data)} 只股票的数据")
        return data_with_target
    
    def _get_all_feature_names(self, integrated_data: Dict[str, Dict[str, Dict[str, Any]]]) -> List[str]:
        """
        获取所有可用特征名称
        
        Args:
            integrated_data: 整合后的数据，键为股票代码，值为日期到特征的映射
            
        Returns:
            所有可用特征的列表
        """
        # 获取所有特征的集合
        all_features = set()
        
        # 遍历所有股票
        for code, data in integrated_data.items():
            # 遍历该股票的所有日期
            for date, daily_data in data.items():
                # 添加当天的所有特征
                all_features.update(daily_data.keys())
        
        # 转为排序的列表返回
        return sorted(list(all_features))
    
    def remove_incomplete_samples(self, data: Dict[str, Dict[str, Dict[str, Any]]]) -> Dict[str, Dict[str, Dict[str, Any]]]:
        """
        移除包含缺失值的样本
        
        Args:
            data: 特征数据，键为股票代码，值为日期到特征的映射
            
        Returns:
            移除缺失值后的数据，键为股票代码，值为日期到特征的映射
        """
        logger.info("开始移除包含缺失值的样本")
        
        cleaned_data = {}
        
        for code, stock_data in data.items():
            cleaned_data[code] = {}
            
            for date, daily_data in stock_data.items():
                # 检查是否有缺失值
                if all(not pd.isna(value) for value in daily_data.values()):
                    cleaned_data[code][date] = daily_data
            
            logger.debug(f"股票 {code}: 原有 {len(stock_data)} 个样本，移除不完整样本后剩余 {len(cleaned_data[code])} 个样本")
        
        logger.info(f"移除不完整样本完成，处理了 {len(data)} 只股票的数据")
        return cleaned_data
    
    def feature_crossing(self, data: Dict[str, Dict[str, Dict[str, float]]], crossing_pairs: List[Tuple[str, str, str]]) -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        创建特征交叉
        
        Args:
            data: 特征数据，键为股票代码，值为日期到特征的映射
            crossing_pairs: 要交叉的特征对列表，每项为(特征1, 特征2, 操作)，操作可为'multiply'或'divide'
            
        Returns:
            包含特征交叉的数据，键为股票代码，值为日期到特征的映射
        """
        logger.info(f"开始创建特征交叉，交叉对数量: {len(crossing_pairs)}")
        
        data_with_crossing = {}
        
        for code, stock_data in data.items():
            data_with_crossing[code] = {}
            
            for date, daily_data in stock_data.items():
                data_with_crossing[code][date] = daily_data.copy()
                
                # 创建特征交叉
                for feature1, feature2, operation in crossing_pairs:
                    if feature1 in daily_data and feature2 in daily_data:
                        value1 = daily_data[feature1]
                        value2 = daily_data[feature2]
                        
                        # 跳过缺失值
                        if pd.isna(value1) or pd.isna(value2):
                            continue
                        
                        # 防止除以零
                        if operation == 'divide' and value2 == 0:
                            continue
                        
                        # 执行交叉操作
                        if operation == 'multiply':
                            data_with_crossing[code][date][f"{feature1}_{operation}_{feature2}"] = value1 * value2
                        elif operation == 'divide':
                            data_with_crossing[code][date][f"{feature1}_{operation}_{feature2}"] = value1 / value2
            
            logger.debug(f"为股票 {code} 创建了特征交叉")
        
        logger.info(f"特征交叉创建完成，处理了 {len(data)} 只股票的数据")
        return data_with_crossing 